package com.xxl.sso.server.util;

import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.engines.SM4Engine;
import org.bouncycastle.crypto.modes.CBCBlockCipher;
import org.bouncycastle.crypto.paddings.PaddedBufferedBlockCipher;
import org.bouncycastle.crypto.params.KeyParameter;
import org.bouncycastle.crypto.params.ParametersWithIV;
import org.bouncycastle.util.encoders.Hex;

import java.nio.charset.StandardCharsets;

public class Sm4Util {
    // Default key and IV (in production, these should be securely managed)
    private static final String DEFAULT_KEY = "0123456789abcdef0123456789abcdef";
    private static final String DEFAULT_IV = "0123456789abcdef0123456789abcdef";

    /**
     * SM4 encryption in CBC mode
     * 
     * @param plainText the text to encrypt
     * @return the hex encoded encrypted data
     */
    public static String encrypt(String plainText) {
        try {
            byte[] key = Hex.decode(DEFAULT_KEY);
            byte[] iv = Hex.decode(DEFAULT_IV);
            
            CipherParameters params = new ParametersWithIV(new KeyParameter(key), iv);
            PaddedBufferedBlockCipher cipher = new PaddedBufferedBlockCipher(new CBCBlockCipher(new SM4Engine()));
            cipher.init(true, params);
            
            byte[] input = plainText.getBytes(StandardCharsets.UTF_8);
            byte[] output = new byte[cipher.getOutputSize(input.length)];
            
            int bytesWritten = cipher.processBytes(input, 0, input.length, output, 0);
            bytesWritten += cipher.doFinal(output, bytesWritten);
            
            return new String(Hex.encode(output, 0, bytesWritten));
        } catch (Exception e) {
            throw new RuntimeException("SM4 encryption failed", e);
        }
    }

    /**
     * SM4 decryption in CBC mode
     * 
     * @param cipherText the hex encoded encrypted data
     * @return the decrypted text
     */
    public static String decrypt(String cipherText) {
        try {
            byte[] key = Hex.decode(DEFAULT_KEY);
            byte[] iv = Hex.decode(DEFAULT_IV);
            
            CipherParameters params = new ParametersWithIV(new KeyParameter(key), iv);
            PaddedBufferedBlockCipher cipher = new PaddedBufferedBlockCipher(new CBCBlockCipher(new SM4Engine()));
            cipher.init(false, params);
            
            byte[] input = Hex.decode(cipherText);
            byte[] output = new byte[cipher.getOutputSize(input.length)];
            
            int bytesWritten = cipher.processBytes(input, 0, input.length, output, 0);
            bytesWritten += cipher.doFinal(output, bytesWritten);
            
            return new String(output, 0, bytesWritten, StandardCharsets.UTF_8);
        } catch (Exception e) {
            throw new RuntimeException("SM4 decryption failed", e);
        }
    }
}